/*
 * lf_decentralized.cc
 *
 * Copyright 2013 Hua Ouyang (houyang@gatech.edu)
 *
 * Decentralized Low Rank Matrix Factorization, solving saddle point of the Lagrangian
 *   L = sum_i |Proj_Omega(A_i - W_i)|_1^2 + 
 *       sum_i<Gamma_i, U_iV_i-W_i> + sum_E <Delta_pn, V_p-V_n>
 *        + 0.5*alpha*(sum_i|U_iV_i-W_i|^2)
 *        + 0.5*rho*(sum_E |V_p-V_n|^2)
 *   where pn \in E means there is an edge from left node p to right node n
 * 
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 * 
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
 * MA 02110-1301, USA.
 * 
 */

#include <armadillo>
#include <boost/archive/text_oarchive.hpp>
#include <boost/archive/text_iarchive.hpp>

#include "helper/helper.h"
#include "algorithm/StopWatch.h"
#include "data/FileAccess.h"
#include <boost/mpi.hpp>
extern boost::mpi::communicator* p_world;
extern po::variables_map vm;
#include "mpi_helper.h"
#include "graph_helper.h"
#include "arma_helper.h"

#define MSG_TAG_A 0
#define MSG_TAG_V 1

class DecentralizedLMF {
public:
  arma::mat A, U, V, W;

  // iteration information
  struct Iteration
  {
    int iter;
    double residual;
    TimeDuration duration, compute, communication;
    Iteration(int iter_, double residual_) : iter(iter_), residual(residual_) {}
  };
  std::vector<Iteration> iterInfo;
protected:
  int m, n, k;
  double alpha, rho;
  int maxiter, interval;
  bool normalizing;
  std::vector<int> neighbors;
  int n_neighbor;
  GraphHelper graph;
  arma::mat Gamma;
  std::vector<arma::mat> V_nb, Delta;
public:
  DecentralizedLMF()
  {
    PO_REQUIRE(vm, "iter");
    PO_REQUIRE(vm, "graph");
    
    PO_VARIABLE_VALUE(rho, vm, "rho", double); // ADMM parameter
    PO_VARIABLE_VALUE(alpha, vm, "alpha", double); // Augemented penalty for Low-Rank MF
    PO_VARIABLE_VALUE(maxiter, vm, "iter", int);
    PO_VARIABLE_VALUE(interval, vm, "interval", int);
    normalizing = vm.count("normalize");

    graph.BuildGraph(neighbors);
    n_neighbor = SZ(neighbors);
  }

  void LoadData() 
  {
    PO_REQUIRE(vm, "data");
    PO_REQUIRE(vm, "k");
    PO_VARIABLE(data, vm, "data", std::string);
    PO_VARIABLE_VALUE(k, vm, "k", int);
    
    using namespace arma;
    std::stringstream ss;
    ss << data << "_" << MPI_RANK;
    A.load(ss.str(), arma_binary);
    m = A.n_rows;
    n = A.n_cols;
  }

  void UpdateU()
  {
    using namespace arma;
    U = trans(solve(V*trans(V), V*trans(W-Gamma/alpha)));
  }

  void UpdateV()
  {
    using namespace arma;
    mpi::request reqs[n_neighbor];
    if (graph.IsLeftSide()) { // left
      // compute V
      V = trans(U) * (alpha * W - Gamma);
      for (int j = 0; j < n_neighbor; j++)
	V += rho * V_nb[j] - Delta[j];
      V = solve(alpha * trans(U) * U + (rho * n_neighbor) * eye(k,k), V);
      if (normalizing) normalize_column(V);

      // send V to neighbors (to right)
      for (int j = 0; j < n_neighbor; j++)
	MPI_ISEND_VERBOSE(reqs[j], neighbors[j], MSG_TAG_V, V, vm, "V["<<MPI_RANK<<"]");
      mpi::wait_all(reqs, reqs+n_neighbor);

      // receive V_neighbors (from right)
      for (int j = 0; j < n_neighbor; j++)
	MPI_IRECV_VERBOSE(reqs[j], neighbors[j], MSG_TAG_V, V_nb[j], vm, "V["<<neighbors[j]<<"]");
      mpi::wait_all(reqs, reqs+n_neighbor);
    }
    else { // right
      // receive V_neighbors (from left)
      for (int i = 0; i < n_neighbor; i++)
	MPI_IRECV_VERBOSE(reqs[i], neighbors[i], MSG_TAG_V, V_nb[i], vm, "V["<<neighbors[i]<<"]");
      mpi::wait_all(reqs, reqs+n_neighbor);
      MSG_VERBOSE(vm, MPI_RANK << " Received V");

      // compute V
      V = trans(U) * (alpha * W - Gamma);
      for (int j = 0; j < n_neighbor; j++)
	V += rho * V_nb[j] + Delta[j];
      V = solve(alpha * trans(U) * U + (rho * n_neighbor) * eye(k,k), V);
      if (normalizing) normalize_column(V);
      
      // send V to neighbors (to left)
      for (int i = 0; i < n_neighbor; i++)
	MPI_ISEND_VERBOSE(reqs[i], neighbors[i], MSG_TAG_V, V, vm, "V["<<MPI_RANK<<"]");
      mpi::wait_all(reqs, reqs+n_neighbor);
    }
  }

  void UpdateW()
  {
    using namespace arma;
    W = U * V - A + Gamma/alpha; 
    soft_threshold(W, 1/alpha);
    W = W + A;
  }

  void UpdateGammaDelta()
  {
    using namespace arma;

    Gamma += alpha*(U*V-W);
    if (graph.IsLeftSide()) { // left
      for (int i = 0; i < n_neighbor; i++)
	Delta[i] += rho*(V-V_nb[i]);
    }
    else { // right
      for (int j = 0; j < n_neighbor; j++)
	Delta[j] += rho*(V_nb[j]-V);
    }
  }

  void SolveLMF() 
  {
    StopWatch watch, clock;
    TimeDuration runtime[2], total_time[2];

    using namespace arma;
    srand(1+MPI_RANK);

    U = randn<mat>(m, k); 
    V = randn<mat>(k, n); 
    W = randn<mat>(m, n); 
    
    Gamma = 10 * randu<mat>(m, n); 
    //Gamma = zeros<mat>(m, n); 

    V_nb = std::vector<mat>(n_neighbor, zeros<mat>(k, n));
    Delta = std::vector<mat>(n_neighbor, zeros<mat>(k, n));

    mpi::request reqs[n_neighbor], reqs_delta[n_neighbor];
    for (int iter=0; iter < maxiter; iter++) {
      // update U in parallel (local)
      MSG_VERBOSE(vm, "--1 it= " << iter << " r=" << MPI_RANK);
      clock.GetDuration();
      UpdateU();
      runtime[0] += clock.GetDuration();

      // update V on left side first, then right side (neighbor send and receive)
      clock.GetDuration();
      UpdateV();
      runtime[1] += clock.GetDuration();
      MSG_VERBOSE(vm, "--2 it= " << iter << " r=" << MPI_RANK);
    
      // update W in parallel (local)
      MSG_VERBOSE(vm, "--3 it= " << iter << " r=" << MPI_RANK);
      clock.GetDuration();
      UpdateW();
      runtime[0] += clock.GetDuration();

      // update dual variables (local)
      clock.GetDuration();
      UpdateGammaDelta();
      runtime[0] += clock.GetDuration();
    
      MSG_VERBOSE(vm, "--4 it= " << iter << " r=" << MPI_RANK);
    
      // print information
      /*
      if (iter % interval == 0 || iter == maxiter-1) {
	double res = norm(A-U*V, "fro"), sum_res;
	res *= res;
	mpi::reduce(*p_world, res, sum_res, std::plus<double>(), 0);
	if (MPI_RANK == 0)
	  MSG("it=" << iter << " |res|=" << sqrt(sum_res)  << " norm(U,fro)=" << norm(U, "fro")
              << " norm(V,fro)=" << norm(V, "fro")
              << " norm(W,fro)=" << norm(W, "fro")
              << " norm(Gam,fro)=" << norm(Gamma, "fro")
	      << " d=" << watch.GetDuration());
      }
      */
      if (iter % interval == 0 || iter == maxiter-1) {
        double res = norm(A-U*V, "fro"), sum_res;
        res *= res;
        mpi::reduce(*p_world, res, sum_res, std::plus<double>(), 0);
        if (MPI_RANK == 0) {
          MSG("it=" << iter << " |res|=" << sqrt(sum_res)  << " norm(U,fro)=" << norm(U, "fro")
              << " norm(V,fro)=" << norm(V, "fro") << " d=" << watch.GetDuration());
	  Iteration info(iter, sqrt(sum_res));
	  info.duration = watch.GetTotalDuration();
	  info.compute = runtime[0];
	  info.communication = runtime[1];
	  iterInfo.push_back(info);
	}
      }
    }
    if (MPI_RANK == 0) {
      MSG(watch.GetTotalDuration());
    }
    mpi::reduce(*p_world, runtime[0], total_time[0], std::plus<TimeDuration>(), 0);
    mpi::reduce(*p_world, runtime[1], total_time[1], std::plus<TimeDuration>(), 0);
    if (MPI_RANK == 0)
      for (int i = 0; i < 2; i++) MSG("clock " << i << " " << total_time[i]);
  }
};


void run_dist_lmf()
{
  if (MPI_RANK == 0) MSG("Decentralized LMF");
  DecentralizedLMF a;
  a.LoadData();
  a.SolveLMF();

  if (MPI_RANK == 0 && vm.count("output")) {
    PO_VARIABLE(output, vm, "output", std::string);
    a.V.save(output, arma::raw_ascii);
  }

  if (MPI_RANK == 0 && vm.count("iterinfo")) {
    PO_VARIABLE(infoFile, vm, "iterinfo", std::string);
    FileAccessWrite(out, infoFile, NULL);
    const std::vector<DecentralizedLMF::Iteration>& info =  a.iterInfo;
    for (unsigned int i = 0; i < info.size(); i++) {
      out << info[i].iter << "\t" << info[i].duration.total_seconds() << "\t" 
	  << info[i].compute.total_seconds() << "\t" 
	  << info[i].communication.total_seconds() << "\t"
	  << info[i].residual << std::endl;
    }
  }
}
